# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import os
import subprocess

import docker
from docker.errors import ContainerError, APIError

from ...utils import logging

EXIT_FAILURE = 1

SPECIAL_PORTS_LIMIT = 1024

logger = logging.get_logger()
logging.set_verbosity_info()

MINDIE_IMAGE_NAME = "registry.modelers.cn/base_image/mindie:1.0.RC3-800I-A2-arm64-OpenMind"
PATH_IN_CONTAINER = "/home/HwHiAiUser/Ascend"

CONTAINER_NAME = "mindie-service"


class DeployMindie:
    def __init__(self, args):
        self.args = args
        self.world_size = args.world_size
        self.port = args.port
        self.model_name_or_path = args.model_name_or_path
        self._set_valid_device()

    @property
    def npu_device_ids(self):
        if isinstance(self.device, str):
            if not all(0 <= int(num) <= 7 for num in self.device.split(",")):
                raise ValueError("Device id must be in range 0 ~ 7 .")
            device = self.device
        else:
            raise TypeError("device data type error, please check data type!")

        logger.info(f"Use device {self.device}")

        return device

    @staticmethod
    def stop_service(remind=False):
        client = docker.from_env()
        containers = client.containers.list(all=True)
        if any(CONTAINER_NAME == container.name for container in containers):
            container = client.containers.get(CONTAINER_NAME)
            container.remove(force=True)
            if remind:
                logger.info("Stop mindie service success.")
        elif remind:
            logger.info("There is no mindie daemon service to stop.")

    def _set_valid_device(self):
        limited_devices = os.getenv("ASCEND_RT_VISIBLE_DEVICES", None)
        if limited_devices is not None:
            self.device = limited_devices
        else:
            if isinstance(self.args.device, int):
                self.device = str(self.args.device)
            else:
                self.device = ",".join(str(x) for x in self.args.device)

    def deploy(self):
        self._check_docker_installed()
        self._init_param()
        self._start_container()

    def generate_npu_device_id(self):
        return ",".join(str(num) for num in range(self.world_size))

    def _check_docker_installed(self):
        try:
            subprocess.run(["docker", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            raise subprocess.CalledProcessError(
                EXIT_FAILURE, "Docker is not installed. Please install it first."
            ) from e
        except FileNotFoundError as e:
            raise FileNotFoundError("Docker is not installed. Please install it first.") from e

    def _init_param(self):
        if len(self.npu_device_ids.split(",")) != self.world_size:
            raise ValueError("Length of devices should be equal with world_size.")

    def _start_container(self):
        self.stop_service()

        client = docker.from_env()
        model_path = os.path.join(PATH_IN_CONTAINER, self.model_name_or_path.rstrip("/").split("/")[-1])
        devices = [f"/dev/davinci{i}" for i in self.npu_device_ids.split(",")]
        devices.extend(["/dev/davinci_manager", "/dev/hisi_hdc", "/dev/devmm_svm"])
        volumes = [
            "/usr/local/Ascend/driver:/usr/local/Ascend/driver",
            "/usr/local/sbin:/usr/local/sbin",
            "/usr/local/dcmi:/usr/local/dcmi",
            "/usr/local/bin/npu-smi:/usr/local/bin/npu-smi",
            f"{self.model_name_or_path}:{model_path}",
        ]

        device_in_container = self.generate_npu_device_id()

        command = (
            f"--model {model_path} "
            f"--port {self.port} "
            f"--npu-device-ids '{device_in_container}' "
            f"--world-size {self.world_size} "
        )
        # since when using the Docker API in Python, there is no progress bar displayed
        self._pull_image()

        try:
            container = client.containers.run(
                image=MINDIE_IMAGE_NAME,
                command=command,
                name=CONTAINER_NAME,
                volumes=volumes,
                ports={f"{self.port}": f"{self.port}"},
                devices=devices,
                detach=True,
            )
            logs = container.logs(stream=True, follow=True)
            for log in logs:
                output = log.decode("utf-8")
                if "Daemon start success" in output:
                    logger.info("Docker run and container init success! ")
                    break
                elif "Fail to get model config path" in output:
                    logger.error("Failed to run MindIE. Please check your model weight path.")
                    break
                elif "Failed to run mindieservice_daemon" in output or "Failed to init endpoint" in output:
                    logger.error("Failed to run MindIE. Please check docker logs.")
                    break

        except (ContainerError, APIError) as err:
            logger.error(f"There is an error during container initialization:{err}")

    def _pull_image(self):
        subprocess.run(["docker", "pull", MINDIE_IMAGE_NAME], check=True, shell=False)
